{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "28a8b793",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/finetuning/embeddings/finetune_embedding.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "551753b7-6cd2-4f81-aec0-da119e4705ad",
   "metadata": {},
   "source": [
    "# Finetune Embeddings\n",
    "\n",
    "In this notebook, we show users how to finetune their own embedding models.\n",
    "\n",
    "We go through three main sections:\n",
    "1. Preparing the data (our `generate_qa_embedding_pairs` function makes this easy)\n",
    "2. Finetuning the model (using our `SentenceTransformersFinetuneEngine`)\n",
    "3. Evaluating the model on a validation knowledge corpus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99afd542-fc47-44ac-aed0-b3684108dba5",
   "metadata": {},
   "source": [
    "## Generate Corpus\n",
    "\n",
    "First, we create the corpus of text chunks by leveraging LlamaIndex to load some financial PDFs, and parsing/chunking into plain text chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9280d438-b6bd-4ccf-a730-7c8bb3ebdbeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "from llama_index import SimpleDirectoryReader\n",
    "from llama_index.node_parser import SentenceSplitter\n",
    "from llama_index.schema import MetadataMode"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "73c42620",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e11b0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/10k/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/lyft_2021.pdf' -O 'data/10k/lyft_2021.pdf'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5e890bc-557b-4d3c-bede-3e80dfeeee18",
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_FILES = [\"./data/10k/lyft_2021.pdf\"]\n",
    "VAL_FILES = [\"./data/10k/uber_2021.pdf\"]\n",
    "\n",
    "TRAIN_CORPUS_FPATH = \"./data/train_corpus.json\"\n",
    "VAL_CORPUS_FPATH = \"./data/val_corpus.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1da871c1-9d58-467a-92fd-06ed3d94534b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_corpus(files, verbose=False):\n",
    "    if verbose:\n",
    "        print(f\"Loading files {files}\")\n",
    "\n",
    "    reader = SimpleDirectoryReader(input_files=files)\n",
    "    docs = reader.load_data()\n",
    "    if verbose:\n",
    "        print(f\"Loaded {len(docs)} docs\")\n",
    "\n",
    "    parser = SentenceSplitter()\n",
    "    nodes = parser.get_nodes_from_documents(docs, show_progress=verbose)\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"Parsed {len(nodes)} nodes\")\n",
    "\n",
    "    return nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53056d8b-3b4c-4364-9b07-a375aa84330b",
   "metadata": {},
   "source": [
    "We do a very naive train/val split by having the Lyft corpus as the train dataset, and the Uber corpus as the val dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3651c77-d085-4fbc-bb34-61f143ad6674",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading files ['./data/10k/lyft_2021.pdf']\n",
      "Loaded 238 docs\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "554a6636780246c8a19d1efe7a6e4786",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Parsing nodes:   0%|          | 0/238 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parsed 344 nodes\n",
      "Loading files ['./data/10k/uber_2021.pdf']\n",
      "Loaded 307 docs\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6748733283a34725ba6365f3c1fb1c1d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Parsing nodes:   0%|          | 0/307 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parsed 410 nodes\n"
     ]
    }
   ],
   "source": [
    "train_nodes = load_corpus(TRAIN_FILES, verbose=True)\n",
    "val_nodes = load_corpus(VAL_FILES, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4482c48-844b-448b-9552-3f38b455645c",
   "metadata": {},
   "source": [
    "### Generate synthetic queries\n",
    "\n",
    "Now, we use an LLM (gpt-3.5-turbo) to generate questions using each text chunk in the corpus as context.\n",
    "\n",
    "Each pair of (generated question, text chunk used as context) becomes a datapoint in the finetuning dataset (either for training or evaluation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "580334ce-ddaa-4cc0-8c3e-7294d11e4d2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.finetuning import (\n",
    "    generate_qa_embedding_pairs,\n",
    "    EmbeddingQAFinetuneDataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "666001e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "OPENAI_API_TOKEN = \"sk-\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_TOKEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef43fe59-a29c-481b-b086-e98e55016d3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 344/344 [12:51<00:00,  2.24s/it]\n",
      "100%|██████████| 410/410 [16:07<00:00,  2.36s/it]\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms import OpenAI\n",
    "\n",
    "\n",
    "train_dataset = generate_qa_embedding_pairs(\n",
    "    llm=OpenAI(model=\"gpt-3.5-turbo\"), nodes=train_nodes\n",
    ")\n",
    "val_dataset = generate_qa_embedding_pairs(\n",
    "    llm=OpenAI(model=\"gpt-3.5-turbo\"), nodes=val_nodes\n",
    ")\n",
    "\n",
    "train_dataset.save_json(\"train_dataset.json\")\n",
    "val_dataset.save_json(\"val_dataset.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "743f163c-25df-4c18-9abe-05052b034d70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Optional] Load\n",
    "train_dataset = EmbeddingQAFinetuneDataset.from_json(\"train_dataset.json\")\n",
    "val_dataset = EmbeddingQAFinetuneDataset.from_json(\"val_dataset.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62368cb8-a303-48b1-8429-5e3655abcc3b",
   "metadata": {},
   "source": [
    "## Run Embedding Finetuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1d08066-5f00-48f1-b12a-e80bc193d4c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.finetuning import SentenceTransformersFinetuneEngine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26625ab5-ddc9-4dbd-9936-39b69c6a7cdc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6efae9c64cdc4a92a248cf1619349958",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       ".gitattributes:   0%|          | 0.00/1.52k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82a58350abe74a59b39686504e56ddb7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "1_Pooling/config.json:   0%|          | 0.00/190 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c938a8515da340fa8567502eb4ab1379",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/90.8k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cec7526cf6d74ab5a90b5a2adecb8dcf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/684 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e8177356c92541939bdc0d7f51a88dd2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config_sentence_transformers.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1155eb2d55b446639814729da89f2a8f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/133M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8db118818c7e4dacb623944d8888e0a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/134M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c7498da6664460ab67ac5fc72fcd565",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "sentence_bert_config.json:   0%|          | 0.00/52.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dfa0168357b74c7f900e49c4cb38b4eb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/125 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39c5ef016f2f4d73bd9d78e081b46f47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/711k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d40419c72754123b95d7a4f3430cac3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/366 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aba92340280a4601a19f4a8707c45fba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.txt:   0%|          | 0.00/232k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e03c93e272574b46a7bb8ca5e389b354",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "modules.json:   0%|          | 0.00/349 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "finetune_engine = SentenceTransformersFinetuneEngine(\n",
    "    train_dataset,\n",
    "    model_id=\"BAAI/bge-small-en\",\n",
    "    model_output_path=\"test_model\",\n",
    "    val_dataset=val_dataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28ad99e6-dd9d-485a-86e9-1845cf51802b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78dab7a09dd640619d80e986baf37249",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b025b81ebe21403498679bf916626ff9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Iteration:   0%|          | 0/69 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e737aae9a5f4459c97df630e63b9c487",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Iteration:   0%|          | 0/69 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "finetune_engine.finetune()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "467a2ba2-e7e6-4025-8887-cac6e7ecb493",
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_model = finetune_engine.get_finetuned_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d16ec01-c29d-4742-aa3c-5ded6ae7c5a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HuggingFaceEmbedding(model_name='test_model', embed_batch_size=10, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x2cc3d5cd0>, tokenizer_name='test_model', max_length=512, pooling=<Pooling.CLS: 'cls'>, normalize=True, query_instruction=None, text_instruction=None, cache_folder=None)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "828dd6fe-9a8a-419b-8663-56d81ce73774",
   "metadata": {},
   "source": [
    "## Evaluate Finetuned Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4a66b83-4cbb-4374-a632-0f1bb2b785ab",
   "metadata": {},
   "source": [
    "In this section, we evaluate 3 different embedding models: \n",
    "1. proprietary OpenAI embedding,\n",
    "2. open source `BAAI/bge-small-en`, and\n",
    "3. our finetuned embedding model.\n",
    "\n",
    "We consider 2 evaluation approaches:\n",
    "1. a simple custom **hit rate** metric\n",
    "2. using `InformationRetrievalEvaluator` from sentence_transformers\n",
    "\n",
    "We show that finetuning on synthetic (LLM-generated) dataset significantly improve upon an opensource embedding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d5176f-1f21-4bcb-adf5-da1c4cccb8d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings import OpenAIEmbedding\n",
    "from llama_index import ServiceContext, VectorStoreIndex\n",
    "from llama_index.schema import TextNode\n",
    "from tqdm.notebook import tqdm\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dda4c2b8-1ad8-420c-83d2-b88e0519895d",
   "metadata": {},
   "source": [
    "### Define eval function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "398c24d3-3d72-4ce8-94a4-2da9c1b2605c",
   "metadata": {},
   "source": [
    "**Option 1**: We use a simple **hit rate** metric for evaluation:\n",
    "* for each (query, relevant_doc) pair,\n",
    "* we retrieve top-k documents with the query,  and \n",
    "* it's a **hit** if the results contain the relevant_doc.\n",
    "\n",
    "This approach is very simple and intuitive, and we can apply it to both the proprietary OpenAI embedding as well as our open source and fine-tuned embedding models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b89401d3-a157-4f96-86d4-212e631a54bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(\n",
    "    dataset,\n",
    "    embed_model,\n",
    "    top_k=5,\n",
    "    verbose=False,\n",
    "):\n",
    "    corpus = dataset.corpus\n",
    "    queries = dataset.queries\n",
    "    relevant_docs = dataset.relevant_docs\n",
    "\n",
    "    service_context = ServiceContext.from_defaults(embed_model=embed_model)\n",
    "    nodes = [TextNode(id_=id_, text=text) for id_, text in corpus.items()]\n",
    "    index = VectorStoreIndex(\n",
    "        nodes, service_context=service_context, show_progress=True\n",
    "    )\n",
    "    retriever = index.as_retriever(similarity_top_k=top_k)\n",
    "\n",
    "    eval_results = []\n",
    "    for query_id, query in tqdm(queries.items()):\n",
    "        retrieved_nodes = retriever.retrieve(query)\n",
    "        retrieved_ids = [node.node.node_id for node in retrieved_nodes]\n",
    "        expected_id = relevant_docs[query_id][0]\n",
    "        is_hit = expected_id in retrieved_ids  # assume 1 relevant doc\n",
    "\n",
    "        eval_result = {\n",
    "            \"is_hit\": is_hit,\n",
    "            \"retrieved\": retrieved_ids,\n",
    "            \"expected\": expected_id,\n",
    "            \"query\": query_id,\n",
    "        }\n",
    "        eval_results.append(eval_result)\n",
    "    return eval_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eb16251-bb45-4de0-b65a-e15aa76e0f1e",
   "metadata": {},
   "source": [
    "**Option 2**: We use the `InformationRetrievalEvaluator` from sentence_transformers.\n",
    "\n",
    "This provides a more comprehensive suite of metrics, but we can only run it against the sentencetransformers compatible models (open source and our finetuned model, *not* the OpenAI embedding model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88e89702-ea35-4c22-99c7-f89a5428ef95",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers.evaluation import InformationRetrievalEvaluator\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "def evaluate_st(\n",
    "    dataset,\n",
    "    model_id,\n",
    "    name,\n",
    "):\n",
    "    corpus = dataset.corpus\n",
    "    queries = dataset.queries\n",
    "    relevant_docs = dataset.relevant_docs\n",
    "\n",
    "    evaluator = InformationRetrievalEvaluator(\n",
    "        queries, corpus, relevant_docs, name=name\n",
    "    )\n",
    "    model = SentenceTransformer(model_id)\n",
    "    output_path = \"results/\"\n",
    "    Path(output_path).mkdir(exist_ok=True, parents=True)\n",
    "    return evaluator(model, output_path=output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af2d33dd-c39f-4c05-8adc-65db12163c88",
   "metadata": {},
   "source": [
    "### Run Evals"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c630aa25-2395-4a8b-83cf-2885fbc862f4",
   "metadata": {},
   "source": [
    "#### OpenAI\n",
    "\n",
    "Note: this might take a few minutes to run since we have to embed the corpus and queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61a0784f-415e-4d3a-8c88-757b28b9e5df",
   "metadata": {},
   "outputs": [],
   "source": [
    "ada = OpenAIEmbedding()\n",
    "ada_val_results = evaluate(val_dataset, ada)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccc73212-fc53-48c1-b347-f5ee3a29ae82",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_ada = pd.DataFrame(ada_val_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25eb61bb-c287-40fe-b3c7-bbfc2d2b1b94",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8779904306220095"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hit_rate_ada = df_ada[\"is_hit\"].mean()\n",
    "hit_rate_ada"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1bd6c62-65a8-4f72-a67c-d0d62c92d7d1",
   "metadata": {},
   "source": [
    "### BAAI/bge-small-en"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24454aeb-9e3e-4954-ab70-647102ed7f82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.011851787567138672,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)ab102/.gitattributes",
       "rate": null,
       "total": 1519,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6e9c5f0555f641caa3a5a5d11cb87583",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)ab102/.gitattributes:   0%|          | 0.00/1.52k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.009984016418457031,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)_Pooling/config.json",
       "rate": null,
       "total": 190,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1fe9a221f8984c818727771d12dfef71",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)_Pooling/config.json:   0%|          | 0.00/190 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.006695985794067383,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)2d2d7ab102/README.md",
       "rate": null,
       "total": 78855,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "619c9cae8bf24987a4d3453aa69d24b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)2d2d7ab102/README.md:   0%|          | 0.00/78.9k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.008507966995239258,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)2d7ab102/config.json",
       "rate": null,
       "total": 684,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "082cfe7c9f3646948886c90f0e1f4258",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)2d7ab102/config.json:   0%|          | 0.00/684 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.011520147323608398,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)ce_transformers.json",
       "rate": null,
       "total": 124,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ff8d7a739fc425abf24076c47c0ab29",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)ce_transformers.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.010421991348266602,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading model.safetensors",
       "rate": null,
       "total": 133466304,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a5344851cb14ed8a5f788cbd74a90d8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading model.safetensors:   0%|          | 0.00/133M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.010251045227050781,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading pytorch_model.bin",
       "rate": null,
       "total": 133508397,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eaa8bdab99244058b1df3eae12a79b20",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/134M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012272834777832031,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)nce_bert_config.json",
       "rate": null,
       "total": 52,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e21b1a35d6c54644be124c357852fedf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)nce_bert_config.json:   0%|          | 0.00/52.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.0071942806243896484,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)cial_tokens_map.json",
       "rate": null,
       "total": 125,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "927efec699ea4c929da7214eb51fc64c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/125 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.010439872741699219,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)ab102/tokenizer.json",
       "rate": null,
       "total": 711396,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c8a00d15090422181a9749e0638e883",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)ab102/tokenizer.json:   0%|          | 0.00/711k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.005940914154052734,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)okenizer_config.json",
       "rate": null,
       "total": 366,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3845bc276c88482ba0e2f2fbe317dd78",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/366 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.00838780403137207,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)2d2d7ab102/vocab.txt",
       "rate": null,
       "total": 231508,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7ceca7b6507e42b1b3da10711b37b7ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)2d2d7ab102/vocab.txt:   0%|          | 0.00/232k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.007328987121582031,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Downloading (…)d7ab102/modules.json",
       "rate": null,
       "total": 229,
       "unit": "B",
       "unit_divisor": 1000,
       "unit_scale": true
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21170e7cf0f9485a9095807a6225aa12",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)d7ab102/modules.json:   0%|          | 0.00/229 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.001703023910522461,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "Generating embeddings",
       "rate": null,
       "total": 418,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3712232b7e064486879945c4d4ac5535",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating embeddings:   0%|          | 0/418 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.0017049312591552734,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 28,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 836,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ba1f47ec020447c59d008493b31e0a57",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/836 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bge = \"local:BAAI/bge-small-en\"\n",
    "bge_val_results = evaluate(val_dataset, bge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2da27e48-1c90-4994-aac4-96b5b1638647",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_bge = pd.DataFrame(bge_val_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ddc4fe0-b240-4c15-9b2d-a4c79f9aaac2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7930622009569378"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hit_rate_bge = df_bge[\"is_hit\"].mean()\n",
    "hit_rate_bge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c16df14-6564-41ec-8816-348675bb0fd4",
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'results/Information-Retrieval_evaluation_bge_results.csv'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[59], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mevaluate_st\u001b[49m\u001b[43m(\u001b[49m\u001b[43mval_dataset\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mBAAI/bge-small-en\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mbge\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[49], line 15\u001b[0m, in \u001b[0;36mevaluate_st\u001b[0;34m(dataset, model_id, name)\u001b[0m\n\u001b[1;32m     13\u001b[0m evaluator \u001b[38;5;241m=\u001b[39m InformationRetrievalEvaluator(queries, corpus, relevant_docs, name\u001b[38;5;241m=\u001b[39mname)\n\u001b[1;32m     14\u001b[0m model \u001b[38;5;241m=\u001b[39m SentenceTransformer(model_id)\n\u001b[0;32m---> 15\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mevaluator\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moutput_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mresults/\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Programming/gpt_index/.venv/lib/python3.10/site-packages/sentence_transformers/evaluation/InformationRetrievalEvaluator.py:104\u001b[0m, in \u001b[0;36mInformationRetrievalEvaluator.__call__\u001b[0;34m(self, model, output_path, epoch, steps, *args, **kwargs)\u001b[0m\n\u001b[1;32m    102\u001b[0m csv_path \u001b[38;5;241m=\u001b[39m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(output_path, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcsv_file)\n\u001b[1;32m    103\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39misfile(csv_path):\n\u001b[0;32m--> 104\u001b[0m     fOut \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mopen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mcsv_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmode\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mw\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mencoding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mutf-8\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    105\u001b[0m     fOut\u001b[38;5;241m.\u001b[39mwrite(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m,\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcsv_headers))\n\u001b[1;32m    106\u001b[0m     fOut\u001b[38;5;241m.\u001b[39mwrite(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'results/Information-Retrieval_evaluation_bge_results.csv'"
     ]
    }
   ],
   "source": [
    "evaluate_st(val_dataset, \"BAAI/bge-small-en\", name=\"bge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fd87550-f547-4b8b-b21a-f72b355e2cd7",
   "metadata": {},
   "source": [
    "### Finetuned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "402dd440-1934-4778-8ff5-28e15cf1f2d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "finetuned = \"local:test_model\"\n",
    "val_results_finetuned = evaluate(val_dataset, finetuned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffd24643-17cb-4773-a535-77f3f8fa2d48",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_finetuned = pd.DataFrame(val_results_finetuned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1dccd1-bbd4-427f-a520-b1011643d83b",
   "metadata": {},
   "outputs": [],
   "source": [
    "hit_rate_finetuned = df_finetuned[\"is_hit\"].mean()\n",
    "hit_rate_finetuned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d8dd38e-f13d-43e1-9802-cc94b854526b",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_st(val_dataset, \"test_model\", name=\"finetuned\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbc290bc-5cc3-4ee4-b8ab-e68371441643",
   "metadata": {},
   "source": [
    "### Summary of Results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f906a11-6a95-4f10-9069-140bf5a56246",
   "metadata": {},
   "source": [
    "#### Hit rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "705fbe3c-2843-4bab-bb5c-16027fc5564b",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_ada[\"model\"] = \"ada\"\n",
    "df_bge[\"model\"] = \"bge\"\n",
    "df_finetuned[\"model\"] = \"fine_tuned\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bebc363c-cd07-4dab-916e-1618d16d1254",
   "metadata": {},
   "source": [
    "We can see that fine-tuning our small open-source embedding model drastically improve its retrieval quality (even approaching the quality of the proprietary OpenAI embedding)!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57f38b4b-1b40-42da-a054-ea9593d3e602",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_all = pd.concat([df_ada, df_bge, df_finetuned])\n",
    "df_all.groupby(\"model\").mean(\"is_hit\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08094d07-2c0a-44ca-ad2f-8d8bf1387ed9",
   "metadata": {},
   "source": [
    "#### InformationRetrievalEvaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27d0444e-a824-42d6-9ddb-4da7179902bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_st_bge = pd.read_csv(\n",
    "    \"results/Information-Retrieval_evaluation_bge_results.csv\"\n",
    ")\n",
    "df_st_finetuned = pd.read_csv(\n",
    "    \"results/Information-Retrieval_evaluation_finetuned_results.csv\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0903ed3-df05-4d98-8b0a-6f352c681735",
   "metadata": {},
   "source": [
    "We can see that embedding finetuning improves metrics consistently across the suite of eval metrics "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81ec1c46-5aa0-4f8a-a0c5-2553e08cceb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_st_bge[\"model\"] = \"bge\"\n",
    "df_st_finetuned[\"model\"] = \"fine_tuned\"\n",
    "df_st_all = pd.concat([df_st_bge, df_st_finetuned])\n",
    "df_st_all = df_st_all.set_index(\"model\")\n",
    "df_st_all"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
